#include <iostream>
#include <chrono>
#include <thread>
#include <vector>
#include <atomic>
#include <algorithm>
#include <cmath>
#include <windows.h>
#include <pdh.h>
#include <pdhmsg.h>
#include <locale>
#include <iomanip> // 用于 std::fixed 和 std::setprecision
#include <string>  // 用于 std::stoi, std::stod, std::wstring
#include <sstream> // 用于 std::wstringstream

#pragma comment(lib, "pdh.lib")

// 用于控制工作线程的全局原子变量
std::atomic<int> g_workTimeMs(250);    // 每个周期的目标工作时长
std::atomic<int> g_sleepTimeMs(250);   // 每个周期的目标休眠时长
std::atomic<bool> g_terminate_threads(false); // 通知线程终止的信号

class CPUMonitor {
private:
    PDH_HQUERY query;
    PDH_HCOUNTER counter;
    int coreIndex; // 新增：用于存储核心索引
    
public:
    // 修改构造函数以接受核心索引
    CPUMonitor(int core = -1) : coreIndex(core) { // core = -1 表示监控总体CPU
        PdhOpenQueryW(nullptr, 0, &query);

        std::wstringstream counterPathStream;
        if (coreIndex == -1) {
            // 监控总体CPU使用率
            counterPathStream << L"\\Processor(_Total)\\% Processor Time";
        } else {
            // 监控特定核心的CPU使用率
            counterPathStream << L"\\Processor(" << coreIndex << L")\\% Processor Time";
        }
        
        PdhAddCounterW(query, counterPathStream.str().c_str(), 0, &counter);
        PdhCollectQueryData(query);
    }

    ~CPUMonitor() {
        // 清理资源
        PdhCloseQuery(query);
    }

    // 获取当前CPU使用率（百分比）
    double getCPUUsage() {
        PDH_FMT_COUNTERVALUE counterVal;
        
        // 收集新数据
        PdhCollectQueryData(query);
        // 获取格式化的计数器值
        PdhGetFormattedCounterValue(counter, PDH_FMT_DOUBLE, nullptr, &counterVal);
        
        return counterVal.doubleValue;
    }
};

// 修改 workerTask 以接受核心索引并设置线程亲和性
void workerTask(int targetCoreIndex) {
    // 设置线程亲和性，确保线程在指定的核心上运行
    if (targetCoreIndex != -1) { // -1 表示不特别指定核心，依赖系统调度
        HANDLE currentThread = GetCurrentThread();
        DWORD_PTR mask = (static_cast<DWORD_PTR>(1) << targetCoreIndex);
        if (!SetThreadAffinityMask(currentThread, mask)) {
            std::cerr << "\n警告: 无法为核心 " << targetCoreIndex << " 设置线程亲和性。错误代码: " << GetLastError() << std::endl;
        }
    }

    while (!g_terminate_threads.load(std::memory_order_relaxed)) {
        int current_work_duration = g_workTimeMs.load(std::memory_order_relaxed);
        
        if (current_work_duration > 0) {
            auto work_period_start_time = std::chrono::high_resolution_clock::now();
            while (std::chrono::duration_cast<std::chrono::milliseconds>(
                       std::chrono::high_resolution_clock::now() - work_period_start_time)
                       .count() < current_work_duration) {
                
                if (g_terminate_threads.load(std::memory_order_relaxed)) break;

                // CPU密集型计算 (来自先前版本)
                volatile double dummy = 0;
                // 这个内层循环是一个工作单元。其强度已在之前的步骤中设置。
                // 大幅减少迭代次数，使单个工作单元更短
                for (int i = 0; i < 50000; ++i) {  // 原为 250000
                    dummy = std::sqrt(dummy + i * 1.1);
                    for (int j = 0; j < 20; ++j) { // 原为 100
                        dummy += std::sin(dummy) + std::cos(dummy) + std::tan(dummy);
                        dummy = std::pow(dummy, 1.01);
                        if (j % 10 == 0) {
                            dummy = std::exp(std::fmod(dummy, 2.0));
                        }
                    }
                    // 在重度循环中定期检查终止信号
                    if (i % 1000 == 0 && g_terminate_threads.load(std::memory_order_relaxed)) { // 调整检查频率
                        break;
                    }
                }
                if (g_terminate_threads.load(std::memory_order_relaxed)) break; 
            }
        }

        if (g_terminate_threads.load(std::memory_order_relaxed)) break;

        int current_sleep_duration = g_sleepTimeMs.load(std::memory_order_relaxed);
        if (current_sleep_duration > 0) {
            // 通过小段休眠实现可中断的休眠
            auto sleep_end_time = std::chrono::steady_clock::now() + std::chrono::milliseconds(current_sleep_duration);
            while (std::chrono::steady_clock::now() < sleep_end_time && !g_terminate_threads.load(std::memory_order_relaxed)) {
                std::this_thread::sleep_for(std::chrono::milliseconds(10)); // 频繁检查终止标志
            }
        }
    }
}

// 获取系统逻辑处理器数量 (从 cpu_core_usage.cpp 引入)
int getProcessorCount() {
    SYSTEM_INFO sysInfo;
    GetSystemInfo(&sysInfo);
    return sysInfo.dwNumberOfProcessors;
}

// 修改 performWorkload 以接受核心索引
void performWorkload(int targetCoreIndex, double targetUsage, int cycleDurationMs, int updateFrequencyMs) {
    CPUMonitor monitor(targetCoreIndex); // 使用核心索引初始化监视器
    
    // 周期内初始的工作/休眠分配
    int workMs = cycleDurationMs / 2;
    
    g_workTimeMs.store(workMs);
    g_sleepTimeMs.store(cycleDurationMs - workMs);
    
    unsigned int num_hw_threads = std::thread::hardware_concurrency();
    if (num_hw_threads == 0) num_hw_threads = 1; // 备用值

    if (targetCoreIndex == -1) {
        std::cout << "开始控制总体CPU使用率在 " << targetUsage << "% 左右..." << std::endl;
    } else {
        std::cout << "开始控制CPU核心 " << targetCoreIndex << " 的使用率在 " << targetUsage << "% 左右..." << std::endl;
    }
    std::cout << "使用 " << num_hw_threads << " 个工作线程。" << std::endl;
    std::cout << "按Ctrl+C退出程序" << std::endl;
    
    while (!g_terminate_threads.load()) {
        double currentUsage = monitor.getCPUUsage();
        
        // 调整逻辑：步长可以是固定的，也可以是 cycleDurationMs 的百分比
        int adjustment_step = std::max(1, cycleDurationMs / 50); // 例如，周期的2%，或至少1毫秒

        if (currentUsage < targetUsage - 2.0) { // CPU过低，增加工作
            workMs += adjustment_step;
        } else if (currentUsage > targetUsage + 2.0) { // CPU过高，减少工作
            workMs -= adjustment_step;
        }
        
        // 确保工作时间在 [0, cycleDurationMs] 范围内
        workMs = std::max(0, std::min(workMs, cycleDurationMs));
        int sleepMs = cycleDurationMs - workMs;
        
        g_workTimeMs.store(workMs);
        g_sleepTimeMs.store(sleepMs);
        
        if (targetCoreIndex == -1) {
            std::cout << "\r总体CPU使用率: " << std::fixed << std::setprecision(2) << currentUsage;
        } else {
            std::cout << "\r核心 " << targetCoreIndex << " CPU使用率: " << std::fixed << std::setprecision(2) << currentUsage;
        }
        std::cout << "%, 目标工作: " << workMs << "ms, 目标睡眠: " << sleepMs 
                  << "ms (周期: " << cycleDurationMs << "ms)     " << std::flush;
        
        // 休眠至更新频率，但定期检查终止标志
        auto sleep_until_next_update = std::chrono::steady_clock::now() + std::chrono::milliseconds(updateFrequencyMs);
        while(std::chrono::steady_clock::now() < sleep_until_next_update && !g_terminate_threads.load(std::memory_order_relaxed)) {
            std::this_thread::sleep_for(std::chrono::milliseconds(10)); 
        }
    }
}

// Ctrl+C 处理函数
BOOL WINAPI ConsoleCtrlHandler(DWORD ctrlType) {
    switch (ctrlType) {
    case CTRL_C_EVENT:
    case CTRL_BREAK_EVENT:
    case CTRL_CLOSE_EVENT:
        if (!g_terminate_threads.load()) {
            std::cout << "\n捕获到Ctrl+C或关闭信号，正在终止线程..." << std::endl;
            g_terminate_threads.store(true);
        }
        // 给线程一些时间反应，在main尝试join它们之前。
        // 返回TRUE表示我们已经处理了信号。程序随后将通过main退出。
        return TRUE;
    default:
        return FALSE; // 传递其他信号
    }
}

int main(int argc, char* argv[]) { // 修改 main 以接受命令行参数
    SetConsoleOutputCP(CP_UTF8); // 设置控制台输出为UTF-8以正确显示中文
    std::cout << std::fixed << std::setprecision(2); // 为double设置输出格式

    // 默认值
    int targetCoreIndex = -1; // -1 表示控制总体CPU使用率
    double targetCpuUsage = 50.0; 
    int processorCount = getProcessorCount();

    // 解析命令行参数
    // 用法: program.exe [core_index] [target_usage_percent]
    if (argc >= 2) {
        if (std::string(argv[1]) == "total") {
            targetCoreIndex = -1;
        } else {
            try {
                targetCoreIndex = std::stoi(argv[1]);
            } catch (const std::exception& e) {
                std::cerr << "无效的核心索引参数: " << argv[1] << ". 请输入一个数字或 'total'." << std::endl;
                std::cerr << "用法: " << argv[0] << " [core_index | total] [target_usage_percent]" << std::endl;
                return 1;
            }
        }
    }

    if (argc >= 3) {
        try {
            targetCpuUsage = std::stod(argv[2]);
            if (targetCpuUsage < 0 || targetCpuUsage > 100) {
                std::cerr << "目标使用率必须在0到100之间。" << std::endl;
                return 1;
            }
        } catch (const std::exception& e) {
            std::cerr << "无效的目标使用率参数: " << argv[2] << std::endl;
            return 1;
        }
    }
    
    if (argc == 1) {
        std::cout << "用法: " << argv[0] << " [core_index | total] [target_usage_percent]" << std::endl;
        std::cout << "未提供参数，将使用默认值：控制总体CPU使用率到 " << targetCpuUsage << "%" << std::endl;
    }


    if (targetCoreIndex != -1) {
        if (targetCoreIndex < 0 || targetCoreIndex >= processorCount) {
            std::cerr << "错误: 核心索引 " << targetCoreIndex << " 超出范围 (0-" << (processorCount - 1) << ")" << std::endl;
            return 1;
        }
        std::cout << "系统共有 " << processorCount << " 个逻辑处理器。目标核心: " << targetCoreIndex << std::endl;
    } else {
        std::cout << "系统共有 " << processorCount << " 个逻辑处理器。目标: 总体CPU使用率" << std::endl;
    }


    unsigned int num_threads_to_spawn = std::thread::hardware_concurrency();
    if (num_threads_to_spawn == 0) {
        num_threads_to_spawn = 4; // 如果检测失败，则使用备用值
    }
    // 如果是针对特定核心，可以考虑减少线程数，例如1或2个。
    // 但当前逻辑是所有线程都绑定到该核心，所以多个线程会竞争该核心。
    // if (targetCoreIndex != -1) {
    //     num_threads_to_spawn = std::min(num_threads_to_spawn, 2u); // 例如，限制为最多2个线程用于单核目标
    // }


    std::vector<std::thread> worker_threads;
    for (unsigned int i = 0; i < num_threads_to_spawn; ++i) {
        // 将目标核心索引传递给每个工作线程
        worker_threads.emplace_back(workerTask, targetCoreIndex);
    }
    
    // 增加周期和更新频率
    int workerCycleDurationMs = 500; 
    int controllerUpdateFrequencyMs = 500; 
    
    if (!SetConsoleCtrlHandler(ConsoleCtrlHandler, TRUE)) {
        std::cerr << "错误: 无法设置控制台Ctrl处理程序。" << std::endl;
        // 不一定是致命错误，但Ctrl+C可能无法正常工作
    }
    
    performWorkload(targetCoreIndex, targetCpuUsage, workerCycleDurationMs, controllerUpdateFrequencyMs);
    
    std::cout << "\n等待所有工作线程退出..." << std::endl;
    // 可选：通知工作线程停止工作并仅休眠，以加快退出速度
    // g_workTimeMs.store(0);
    // g_sleepTimeMs.store(workerCycleDurationMs);

    for (auto& t : worker_threads) {
        if (t.joinable()) {
            t.join();
        }
    }
    
    std::cout << "所有工作线程已成功退出。程序结束。" << std::endl;
    return 0;
}
